"""
DaoMath Quantum Mechanics (道数量子力学)
从 ±0 先天均衡构造量子理论

量子力学的 ±0 诠释:
    波 ↔ 粒 (对偶)
    位置 ↔ 动量 (对偶)
    能量 ↔ 时间 (对偶)
    测量前 ↔ 测量后 (对偶)
    
理论层级:
    ±0 (先天均衡)
     ↓
    波函数 (State Vector)
     ↓
    算符 (Operators)
     ↓
    Schrödinger 方程
     ↓
    测量理论
     ↓
    纠缠与非定域性
     ↓
    路径积分
"""

from __future__ import annotations
from typing import List, Tuple, Callable, Optional, Dict, Any, Union
from dataclasses import dataclass
from abc import ABC, abstractmethod
import numpy as np
import scipy.linalg as la
from scipy.integrate import odeint, solve_ivp
from scipy.special import hermite, genlaguerre, sph_harm, factorial
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from core.axioms import InnateBalance, DualPair


# ============================================================================
# 物理常数
# ============================================================================

class QuantumConstants:
    """量子力学常数"""
    HBAR = 1.054571817e-34      # ℏ (J·s)
    M_E = 9.1093837015e-31      # 电子质量 (kg)
    E_CHARGE = 1.602176634e-19  # 基本电荷 (C)
    EPSILON_0 = 8.8541878128e-12  # 真空介电常数 (F/m)
    
    # 原子单位制
    HBAR_AU = 1.0
    M_E_AU = 1.0
    E_CHARGE_AU = 1.0


# ============================================================================
# Level 1: 量子态与希尔伯特空间
# ============================================================================

class QuantumState(ABC):
    """
    量子态抽象基类
    
    |ψ⟩ ∈ ℋ (Hilbert 空间)
    
    基于 ±0 的理解:
    - 叠加态：正负振幅的相干叠加
    - 归一化：总概率为 1 (守恒)
    """
    
    @abstractmethod
    def inner_product(self, other: QuantumState) -> complex:
        """内积 ⟨ψ|φ⟩"""
        pass
    
    @abstractmethod
    def norm(self) -> float:
        """模 ||ψ||"""
        pass
    
    @abstractmethod
    def normalize(self) -> QuantumState:
        """归一化"""
        pass
    
    def is_normalized(self, tol: float = 1e-10) -> bool:
        """检查归一化"""
        return abs(self.norm() - 1.0) < tol
    
    def orthogonal_to(self, other: QuantumState, tol: float = 1e-10) -> bool:
        """检查正交性"""
        return abs(self.inner_product(other)) < tol


class WaveFunction(QuantumState):
    """
    波函数（位置表象）
    
    ψ(x, t): 空间位置 → 复振幅
    
    |ψ(x)|² = 概率密度
    """
    
    def __init__(self, 
                 function: Callable[[np.ndarray, float], complex],
                 domain: Tuple[np.ndarray, np.ndarray],
                 dimension: int = 1):
        """
        Args:
            function: (x, t) → ψ(x, t)
            domain: (x_min, x_max)
            dimension: 空间维度
        """
        self.psi = function
        self.x_min, self.x_max = domain
        self.dim = dimension
        self.time = 0.0
    
    def __call__(self, x: np.ndarray, t: float = 0.0) -> complex:
        """计算波函数值"""
        return self.psi(x, t)
    
    def inner_product(self, other: WaveFunction, 
                     num_points: int = 1000) -> complex:
        """
        内积: ⟨ψ|φ⟩ = ∫ ψ*(x) φ(x) dx
        """
        x = np.linspace(self.x_min[0], self.x_max[0], num_points)
        dx = (self.x_max[0] - self.x_min[0]) / num_points
        
        psi_vals = np.array([self(np.array([xi]), self.time) for xi in x])
        phi_vals = np.array([other(np.array([xi]), other.time) for xi in x])
        
        return np.sum(np.conj(psi_vals) * phi_vals) * dx
    
    def norm(self) -> float:
        """模 ||ψ|| = √⟨ψ|ψ⟩"""
        return np.sqrt(abs(self.inner_product(self)))
    
    def normalize(self) -> WaveFunction:
        """归一化波函数"""
        N = self.norm()
        
        def normalized_psi(x: np.ndarray, t: float) -> complex:
            return self.psi(x, t) / N
        
        return WaveFunction(normalized_psi, (self.x_min, self.x_max), self.dim)
    
    def probability_density(self, x: np.ndarray, t: float = 0.0) -> float:
        """
        概率密度: ρ(x) = |ψ(x)|²
        """
        psi_val = self(x, t)
        return abs(psi_val)**2
    
    def expectation_value(self, 
                         operator: Callable[[np.ndarray], complex],
                         num_points: int = 1000) -> complex:
        """
        期望值: ⟨Ô⟩ = ∫ ψ*(x) Ô ψ(x) dx
        """
        x_vals = np.linspace(self.x_min[0], self.x_max[0], num_points)
        dx = (self.x_max[0] - self.x_min[0]) / num_points
        
        expectation = 0.0 + 0.0j
        for x in x_vals:
            psi = self(np.array([x]), self.time)
            O_psi = operator(np.array([x]))
            expectation += np.conj(psi) * O_psi * psi
        
        return expectation * dx
    
    def position_expectation(self, num_points: int = 1000) -> float:
        """位置期望值 ⟨x⟩"""
        x_vals = np.linspace(self.x_min[0], self.x_max[0], num_points)
        dx = (self.x_max[0] - self.x_min[0]) / num_points
        
        x_exp = 0.0
        for x in x_vals:
            rho = self.probability_density(np.array([x]), self.time)
            x_exp += x * rho
        
        return x_exp * dx
    
    def momentum_expectation(self, 
                            hbar: float = QuantumConstants.HBAR_AU,
                            num_points: int = 1000) -> float:
        """
        动量期望值 ⟨p⟩ = -iℏ ∫ ψ* ∂ψ/∂x dx
        """
        x_vals = np.linspace(self.x_min[0], self.x_max[0], num_points)
        dx = (self.x_max[0] - self.x_min[0]) / num_points
        
        p_exp = 0.0 + 0.0j
        for i, x in enumerate(x_vals[:-1]):
            psi = self(np.array([x]), self.time)
            psi_next = self(np.array([x_vals[i+1]]), self.time)
            
            dpsi_dx = (psi_next - psi) / dx
            p_exp += np.conj(psi) * (-1j * hbar * dpsi_dx)
        
        return np.real(p_exp * dx)
    
    @staticmethod
    def gaussian_wavepacket(x0: float, 
                           k0: float, 
                           sigma: float,
                           domain: Tuple[np.ndarray, np.ndarray]) -> WaveFunction:
        """
        高斯波包
        
        ψ(x) = (2πσ²)^(-1/4) exp(-(x-x₀)²/4σ²) exp(ik₀x)
        
        体现 ±0: 位置与动量的不确定性均衡
        """
        def psi(x: np.ndarray, t: float) -> complex:
            norm = (2 * np.pi * sigma**2)**(-0.25)
            spatial = np.exp(-(x[0] - x0)**2 / (4 * sigma**2))
            momentum = np.exp(1j * k0 * x[0])
            return norm * spatial * momentum
        
        return WaveFunction(psi, domain)


class StateVector(QuantumState):
    """
    态矢量（离散表象）
    
    |ψ⟩ = Σ_n c_n |n⟩
    
    有限维希尔伯特空间
    """
    
    def __init__(self, coefficients: np.ndarray):
        """
        Args:
            coefficients: 基态展开系数 [c_0, c_1, ..., c_n]
        """
        self.coeffs = np.array(coefficients, dtype=complex)
        self.dim = len(coefficients)
    
    def inner_product(self, other: StateVector) -> complex:
        """内积: ⟨ψ|φ⟩ = Σ_n ψ_n* φ_n"""
        if self.dim != other.dim:
            raise ValueError("Dimensions must match")
        
        return np.vdot(self.coeffs, other.coeffs)
    
    def norm(self) -> float:
        """模"""
        return np.sqrt(abs(self.inner_product(self)))
    
    def normalize(self) -> StateVector:
        """归一化"""
        N = self.norm()
        return StateVector(self.coeffs / N)
    
    def to_density_matrix(self) -> DensityMatrix:
        """
        转换为密度矩阵
        
        ρ = |ψ⟩⟨ψ|
        """
        return DensityMatrix(np.outer(self.coeffs, np.conj(self.coeffs)))
    
    @staticmethod
    def basis_state(n: int, dim: int) -> StateVector:
        """
        基态 |n⟩
        """
        coeffs = np.zeros(dim, dtype=complex)
        coeffs[n] = 1.0
        return StateVector(coeffs)
    
    @staticmethod
    def superposition(states: List[StateVector], 
                     amplitudes: List[complex]) -> StateVector:
        """
        叠加态: |ψ⟩ = Σ_i c_i |ψ_i⟩
        
        体现 ±0: 多个态的相干叠加
        """
        if len(states) != len(amplitudes):
            raise ValueError("Number of states and amplitudes must match")
        
        dim = states[0].dim
        result = np.zeros(dim, dtype=complex)
        
        for state, amp in zip(states, amplitudes):
            result += amp * state.coeffs
        
        return StateVector(result)


class DensityMatrix:
    """
    密度矩阵
    
    ρ = Σ_i p_i |ψ_i⟩⟨ψ_i|
    
    描述混合态
    """
    
    def __init__(self, matrix: np.ndarray):
        """
        Args:
            matrix: 密度矩阵 (Hermitian, 半正定, Tr(ρ)=1)
        """
        self.rho = np.array(matrix, dtype=complex)
        self.dim = matrix.shape[0]
    
    def trace(self) -> float:
        """迹 Tr(ρ)"""
        return np.real(np.trace(self.rho))
    
    def is_normalized(self, tol: float = 1e-10) -> bool:
        """检查归一化 Tr(ρ) = 1"""
        return abs(self.trace() - 1.0) < tol
    
    def is_pure(self, tol: float = 1e-10) -> bool:
        """
        纯态检验: Tr(ρ²) = 1
        """
        rho2 = self.rho @ self.rho
        return abs(np.trace(rho2) - 1.0) < tol
    
    def von_neumann_entropy(self) -> float:
        """
        von Neumann 熵
        
        S = -Tr(ρ ln ρ)
        
        体现 ±0: 纯态熵=0, 最大混合态熵最大
        """
        eigenvalues = np.linalg.eigvalsh(self.rho)
        eigenvalues = eigenvalues[eigenvalues > 1e-15]  # 避免 log(0)
        
        return -np.sum(eigenvalues * np.log(eigenvalues))
    
    def expectation_value(self, operator: np.ndarray) -> complex:
        """
        期望值: ⟨Ô⟩ = Tr(ρÔ)
        """
        return np.trace(self.rho @ operator)
    
    def partial_trace(self, subsystem: int, dims: Tuple[int, int]) -> DensityMatrix:
        """
        偏迹（约化密度矩阵）
        
        对于复合系统 A⊗B, 求子系统密度矩阵
        """
        dim_A, dim_B = dims
        
        if subsystem == 0:  # 保留系统 A
            rho_A = np.zeros((dim_A, dim_A), dtype=complex)
            for i in range(dim_A):
                for j in range(dim_A):
                    for k in range(dim_B):
                        idx1 = i * dim_B + k
                        idx2 = j * dim_B + k
                        rho_A[i, j] += self.rho[idx1, idx2]
            return DensityMatrix(rho_A)
        
        else:  # 保留系统 B
            rho_B = np.zeros((dim_B, dim_B), dtype=complex)
            for i in range(dim_B):
                for j in range(dim_B):
                    for k in range(dim_A):
                        idx1 = k * dim_B + i
                        idx2 = k * dim_B + j
                        rho_B[i, j] += self.rho[idx1, idx2]
            return DensityMatrix(rho_B)


# ============================================================================
# Level 2: 量子算符
# ============================================================================

class QuantumOperator(ABC):
    """
    量子算符抽象基类
    
    Ô: ℋ → ℋ
    
    基于 ±0 的理解:
    - Hermitian 算符：本征值实数（可观测量）
    - Unitary 算符：保模（演化算符）
    """
    
    @abstractmethod
    def apply(self, state: QuantumState) -> QuantumState:
        """作用于态"""
        pass
    
    @abstractmethod
    def matrix_element(self, bra: QuantumState, ket: QuantumState) -> complex:
        """矩阵元 ⟨ψ|Ô|φ⟩"""
        pass


class HermitianOperator(QuantumOperator):
    """
    Hermitian 算符（可观测量）
    
    Ô† = Ô
    
    本征值实数，本征态正交
    """
    
    def __init__(self, matrix: np.ndarray):
        """
        Args:
            matrix: Hermitian 矩阵
        """
        if not np.allclose(matrix, matrix.conj().T):
            raise ValueError("Matrix must be Hermitian")
        
        self.matrix = matrix
        self.dim = matrix.shape[0]
        
        # 本征值分解
        self.eigenvalues, self.eigenvectors = np.linalg.eigh(matrix)
    
    def apply(self, state: StateVector) -> StateVector:
        """Ô|ψ⟩"""
        new_coeffs = self.matrix @ state.coeffs
        return StateVector(new_coeffs)
    
    def matrix_element(self, bra: StateVector, ket: StateVector) -> complex:
        """⟨ψ|Ô|φ⟩"""
        return np.vdot(bra.coeffs, self.matrix @ ket.coeffs)
    
    def expectation_value(self, state: StateVector) -> float:
        """
        期望值: ⟨Ô⟩ = ⟨ψ|Ô|ψ⟩
        """
        return np.real(self.matrix_element(state, state))
    
    def uncertainty(self, state: StateVector) -> float:
        """
        不确定度: ΔÔ = √(⟨Ô²⟩ - ⟨Ô⟩²)
        """
        O_avg = self.expectation_value(state)
        O2_avg = self.expectation_value_squared(state)
        
        variance = O2_avg - O_avg**2
        return np.sqrt(max(0, variance))
    
    def expectation_value_squared(self, state: StateVector) -> float:
        """⟨Ô²⟩"""
        O2 = HermitianOperator(self.matrix @ self.matrix)
        return O2.expectation_value(state)
    
    def commutator(self, other: HermitianOperator) -> np.ndarray:
        """
        对易子: [Ô₁, Ô₂] = Ô₁Ô₂ - Ô₂Ô₁
        
        体现 ±0: 正向减负向
        """
        return self.matrix @ other.matrix - other.matrix @ self.matrix
    
    def uncertainty_relation(self, 
                            other: HermitianOperator,
                            state: StateVector) -> Tuple[float, float, float]:
        """
        不确定性关系
        
        ΔÔ₁ · ΔÔ₂ ≥ |⟨[Ô₁,Ô₂]⟩| / 2
        
        Returns:
            (ΔÔ₁, ΔÔ₂, 下界)
        """
        dO1 = self.uncertainty(state)
        dO2 = other.uncertainty(state)
        
        comm = self.commutator(other)
        comm_op = HermitianOperator((comm - comm.conj().T) / (2j))
        comm_expectation = abs(comm_op.expectation_value(state))
        
        bound = comm_expectation / 2
        
        return (dO1, dO2, bound)


class PositionOperator(HermitianOperator):
    """
    位置算符 x̂
    
    在位置表象: x̂ ψ(x) = x ψ(x)
    """
    
    def __init__(self, dim: int, domain: Tuple[float, float]):
        """
        Args:
            dim: 离散化维度
            domain: (x_min, x_max)
        """
        x_min, x_max = domain
        x_vals = np.linspace(x_min, x_max, dim)
        
        # 对角矩阵
        matrix = np.diag(x_vals)
        super().__init__(matrix)


class MomentumOperator(HermitianOperator):
    """
    动量算符 p̂
    
    在位置表象: p̂ = -iℏ ∂/∂x
    """
    
    def __init__(self, 
                 dim: int,
                 dx: float,
                 hbar: float = QuantumConstants.HBAR_AU):
        """
        Args:
            dim: 维度
            dx: 空间步长
            hbar: ℏ
        """
        # 有限差分近似
        matrix = np.zeros((dim, dim), dtype=complex)
        
        for i in range(dim):
            if i > 0:
                matrix[i, i-1] = -1
            if i < dim - 1:
                matrix[i, i+1] = 1
        
        matrix *= -1j * hbar / (2 * dx)
        
        # 确保 Hermitian
        matrix = (matrix + matrix.conj().T) / 2
        
        super().__init__(matrix)


class HamiltonianOperator(HermitianOperator):
    """
    Hamiltonian 算符 Ĥ
    
    能量算符，生成时间演化
    """
    
    def __init__(self, matrix: np.ndarray):
        super().__init__(matrix)
    
    def time_evolution_operator(self, 
                                time: float,
                                hbar: float = QuantumConstants.HBAR_AU) -> UnitaryOperator:
        """
        时间演化算符
        
        Û(t) = exp(-iĤt/ℏ)
        """
        U_matrix = la.expm(-1j * self.matrix * time / hbar)
        return UnitaryOperator(U_matrix)
    
    def energy_eigenstate(self, n: int) -> Tuple[float, StateVector]:
        """
        第 n 个能量本征态
        
        Returns:
            (本征能量, 本征态)
        """
        energy = self.eigenvalues[n]
        state = StateVector(self.eigenvectors[:, n])
        
        return (energy, state)
    
    @staticmethod
    def harmonic_oscillator(dim: int, 
                           omega: float = 1.0,
                           hbar: float = QuantumConstants.HBAR_AU,
                           mass: float = 1.0) -> HamiltonianOperator:
        """
        谐振子 Hamiltonian
        
        Ĥ = p²/2m + ½mω²x²
        
        体现 ±0: 动能(+) 与势能(+) 的均衡
        """
        # 使用产生湮灭算符构造
        # Ĥ = ℏω(â†â + ½)
        
        matrix = np.zeros((dim, dim))
        
        for n in range(dim):
            matrix[n, n] = hbar * omega * (n + 0.5)
        
        return HamiltonianOperator(matrix)
    
    @staticmethod
    def hydrogen_atom(n_max: int,
                     l: int,
                     Z: int = 1) -> HamiltonianOperator:
        """
        氢原子 Hamiltonian（简化为径向部分）
        
        E_n = -Z² Ry / n²
        
        其中 Ry = 13.6 eV (Rydberg 常数)
        """
        dim = n_max - l
        matrix = np.zeros((dim, dim))
        
        Ry = 13.6  # eV
        
        for i in range(dim):
            n = l + i + 1
            matrix[i, i] = -Z**2 * Ry / n**2
        
        return HamiltonianOperator(matrix)


class UnitaryOperator(QuantumOperator):
    """
    Unitary 算符
    
    Û†Û = ÛÛ† = Î
    
    保持态矢量归一化（时间演化、旋转等）
    """
    
    def __init__(self, matrix: np.ndarray):
        """
        Args:
            matrix: Unitary 矩阵
        """
        # 验证 Unitary 性
        identity = np.eye(matrix.shape[0])
        if not np.allclose(matrix @ matrix.conj().T, identity):
            raise ValueError("Matrix must be unitary")
        
        self.matrix = matrix
        self.dim = matrix.shape[0]
    
    def apply(self, state: StateVector) -> StateVector:
        """Û|ψ⟩"""
        new_coeffs = self.matrix @ state.coeffs
        return StateVector(new_coeffs)
    
    def matrix_element(self, bra: StateVector, ket: StateVector) -> complex:
        """⟨ψ|Û|φ⟩"""
        return np.vdot(bra.coeffs, self.matrix @ ket.coeffs)
    
    def inverse(self) -> UnitaryOperator:
        """逆算符 Û⁻¹ = Û†"""
        return UnitaryOperator(self.matrix.conj().T)


# ============================================================================
# Level 3: Schrödinger 方程
# ============================================================================

class SchrodingerEquation:
    """
    Schrödinger 方程
    
    含时: iℏ ∂ψ/∂t = Ĥψ
    定态: Ĥψ = Eψ
    
    基于 ±0 的理解:
    - 能量守恒（Hamiltonian 生成时间演化）
    - 概率守恒（Unitary 演化）
    """
    
    def __init__(self, 
                 hamiltonian: HamiltonianOperator,
                 hbar: float = QuantumConstants.HBAR_AU):
        """
        Args:
            hamiltonian: Hamiltonian 算符
            hbar: ℏ
        """
        self.H = hamiltonian
        self.hbar = hbar
    
    def time_evolution(self,
                      initial_state: StateVector,
                      time_span: Tuple[float, float],
                      num_points: int = 100) -> Tuple[np.ndarray, List[StateVector]]:
        """
        求解含时 Schrödinger 方程
        
        Returns:
            (时间点, 态列表)
        """
        t_vals = np.linspace(time_span[0], time_span[1], num_points)
        states = []
        
        for t in t_vals:
            U = self.H.time_evolution_operator(t - time_span[0], self.hbar)
            state_t = U.apply(initial_state)
            states.append(state_t)
        
        return (t_vals, states)
    
    def eigenstate_expansion(self,
                            initial_state: StateVector,
                            time: float) -> StateVector:
        """
        本征态展开法求解
        
        ψ(t) = Σ_n c_n e^(-iE_nt/ℏ) |n⟩
        """
        # 展开到本征态
        coeffs = np.zeros(self.H.dim, dtype=complex)
        
        for n in range(self.H.dim):
            E_n, phi_n = self.H.energy_eigenstate(n)
            c_n = phi_n.inner_product(initial_state)
            
            # 时间演化因子
            phase = np.exp(-1j * E_n * time / self.hbar)
            coeffs += c_n * phase * phi_n.coeffs
        
        return StateVector(coeffs)
    
    def stationary_states(self) -> List[Tuple[float, StateVector]]:
        """
        求解定态 Schrödinger 方程
        
        Returns:
            [(E_0, |0⟩), (E_1, |1⟩), ...]
        """
        states = []
        for n in range(self.H.dim):
            states.append(self.H.energy_eigenstate(n))
        
        return states


# ============================================================================
# Level 4: 量子测量理论
# ============================================================================

class Measurement:
    """
    量子测量
    
    测量：态的塌缩
    
    基于 ±0 的理解:
    - 测量前：叠加态（±均衡）
    - 测量后：本征态（确定值）
    - 概率：Born 规则
    """
    
    def __init__(self, observable: HermitianOperator):
        """
        Args:
            observable: 可观测量算符
        """
        self.observable = observable
    
    def measurement_probability(self, 
                               state: StateVector,
                               eigenstate_index: int) -> float:
        """
        测得第 n 个本征值的概率
        
        P_n = |⟨n|ψ⟩|²
        
        Born 规则
        """
        eigenstate = StateVector(self.observable.eigenvectors[:, eigenstate_index])
        amplitude = eigenstate.inner_product(state)
        
        return abs(amplitude)**2
    
    def measure(self, state: StateVector) -> Tuple[float, StateVector]:
        """
        执行测量
        
        Returns:
            (测量结果, 塌缩后态)
        """
        # 计算所有本征态的概率
        probabilities = [
            self.measurement_probability(state, n)
            for n in range(self.observable.dim)
        ]
        
        # 随机选择（根据概率）
        n = np.random.choice(self.observable.dim, p=probabilities)
        
        # 测量结果
        result = self.observable.eigenvalues[n]
        
        # 塌缩到本征态
        collapsed_state = StateVector(self.observable.eigenvectors[:, n])
        
        return (result, collapsed_state)
    
    def expectation_value(self, state: StateVector) -> float:
        """
        期望值: ⟨Ô⟩ = Σ_n P_n λ_n
        """
        return self.observable.expectation_value(state)
    
    def variance(self, state: StateVector) -> float:
        """
        方差: Var(Ô) = ⟨Ô²⟩ - ⟨Ô⟩²
        """
        return self.observable.uncertainty(state)**2


class POVMMeasurement:
    """
    广义测量（POVM）
    
    Positive Operator-Valued Measure
    
    比 von Neumann 测量更一般
    """
    
    def __init__(self, elements: List[np.ndarray]):
        """
        Args:
            elements: POVM 元素 {E_m}, 满足 Σ_m E_m = I
        """
        self.elements = elements
        self.num_outcomes = len(elements)
        
        # 验证完备性
        total = sum(elements)
        dim = elements[0].shape[0]
        if not np.allclose(total, np.eye(dim)):
            raise ValueError("POVM elements must sum to identity")
    
    def measurement_probability(self, 
                               state: StateVector,
                               outcome: int) -> float:
        """
        测量结果 m 的概率
        
        P(m) = ⟨ψ|E_m|ψ⟩
        """
        E_m = self.elements[outcome]
        return np.real(np.vdot(state.coeffs, E_m @ state.coeffs))
    
    def measure(self, state: StateVector) -> int:
        """
        执行 POVM 测量
        
        Returns:
            测量结果索引
        """
        probabilities = [
            self.measurement_probability(state, m)
            for m in range(self.num_outcomes)
        ]
        
        return np.random.choice(self.num_outcomes, p=probabilities)


# ============================================================================
# Level 5: 量子纠缠与非定域性
# ============================================================================

class EntangledState:
    """
    纠缠态
    
    不能写成直积形式的态
    
    |ψ⟩_{AB} ≠ |ψ_A⟩ ⊗ |ψ_B⟩
    
    体现 ±0: 子系统间的非局域关联
    """
    
    def __init__(self, 
                 state: StateVector,
                 subsystem_dims: Tuple[int, int]):
        """
        Args:
            state: 复合系统态矢量
            subsystem_dims: (dim_A, dim_B)
        """
        self.state = state
        self.dim_A, self.dim_B = subsystem_dims
        
        if state.dim != self.dim_A * self.dim_B:
            raise ValueError("State dimension must match subsystem dimensions")
    
    def reduced_density_matrix(self, subsystem: int) -> DensityMatrix:
        """
        约化密度矩阵
        
        ρ_A = Tr_B(|ψ⟩⟨ψ|)
        """
        full_rho = self.state.to_density_matrix()
        return full_rho.partial_trace(subsystem, (self.dim_A, self.dim_B))
    
    def entanglement_entropy(self, subsystem: int = 0) -> float:
        """
        纠缠熵
        
        S_E = -Tr(ρ_A ln ρ_A)
        
        体现纠缠程度
        """
        rho_reduced = self.reduced_density_matrix(subsystem)
        return rho_reduced.von_neumann_entropy()
    
    def is_entangled(self, tol: float = 1e-10) -> bool:
        """
        判断是否纠缠
        
        方法：检查约化密度矩阵是否为纯态
        """
        rho_A = self.reduced_density_matrix(0)
        return not rho_A.is_pure(tol)
    
    def schmidt_decomposition(self) -> Tuple[np.ndarray, List[StateVector], List[StateVector]]:
        """
        Schmidt 分解
        
        |ψ⟩ = Σ_i √λ_i |i_A⟩|i_B⟩
        
        Returns:
            (Schmidt 系数, A 基态, B 基态)
        """
        # 重塑为矩阵
        psi_matrix = self.state.coeffs.reshape(self.dim_A, self.dim_B)
        
        # SVD 分解
        U, s, Vh = np.linalg.svd(psi_matrix)
        
        # Schmidt 系数
        schmidt_coeffs = s
        
        # 基态
        states_A = [StateVector(U[:, i]) for i in range(len(s))]
        states_B = [StateVector(Vh[i, :]) for i in range(len(s))]
        
        return (schmidt_coeffs, states_A, states_B)
    
    def schmidt_number(self) -> int:
        """
        Schmidt 数（非零 Schmidt 系数个数）
        
        纠缠度量
        """
        coeffs, _, _ = self.schmidt_decomposition()
        return np.sum(coeffs > 1e-10)
    
    @staticmethod
    def bell_state(which: str = 'phi_plus') -> EntangledState:
        """
        Bell 态（最大纠缠态）
        
        |Φ⁺⟩ = (|00⟩ + |11⟩)/√2
        |Φ⁻⟩ = (|00⟩ - |11⟩)/√2
        |Ψ⁺⟩ = (|01⟩ + |10⟩)/√2
        |Ψ⁻⟩ = (|01⟩ - |10⟩)/√2
        
        体现 ±0: 完美的正负关联
        """
        if which == 'phi_plus':
            coeffs = np.array([1, 0, 0, 1]) / np.sqrt(2)
        elif which == 'phi_minus':
            coeffs = np.array([1, 0, 0, -1]) / np.sqrt(2)
        elif which == 'psi_plus':
            coeffs = np.array([0, 1, 1, 0]) / np.sqrt(2)
        elif which == 'psi_minus':
            coeffs = np.array([0, 1, -1, 0]) / np.sqrt(2)
        else:
            raise ValueError(f"Unknown Bell state: {which}")
        
        state = StateVector(coeffs)
        return EntangledState(state, (2, 2))
    
    @staticmethod
    def ghz_state(n_qubits: int) -> EntangledState:
        """
        GHZ 态（多粒子纠缠）
        
        |GHZ⟩ = (|0...0⟩ + |1...1⟩)/√2
        """
        dim = 2**n_qubits
        coeffs = np.zeros(dim, dtype=complex)
        coeffs[0] = 1 / np.sqrt(2)
        coeffs[-1] = 1 / np.sqrt(2)
        
        state = StateVector(coeffs)
        return EntangledState(state, (2, 2**(n_qubits-1)))


class BellInequality:
    """
    Bell 不等式
    
    检验量子力学的非定域性
    
    CHSH 不等式: |S| ≤ 2 (经典)
                |S| ≤ 2√2 (量子)
    """
    
    @staticmethod
    def chsh_parameter(state: EntangledState,
                       observables_A: Tuple[np.ndarray, np.ndarray],
                       observables_B: Tuple[np.ndarray, np.ndarray]) -> float:
        """
        CHSH 参数
        
        S = E(A₁B₁) + E(A₁B₂) + E(A₂B₁) - E(A₂B₂)
        
        Args:
            observables_A: (A₁, A₂) Alice 的观测算符
            observables_B: (B₁, B₂) Bob 的观测算符
        
        Returns:
            S 值
        """
        A1, A2 = observables_A
        B1, B2 = observables_B
        
        # 构造联合测量算符
        def joint_observable(A, B):
            return np.kron(A, B)
        
        # 计算关联函数
        def correlation(O_joint):
            rho = state.state.to_density_matrix()
            return np.real(rho.expectation_value(O_joint))
        
        E_A1B1 = correlation(joint_observable(A1, B1))
        E_A1B2 = correlation(joint_observable(A1, B2))
        E_A2B1 = correlation(joint_observable(A2, B1))
        E_A2B2 = correlation(joint_observable(A2, B2))
        
        S = E_A1B1 + E_A1B2 + E_A2B1 - E_A2B2
        
        return S
    
    @staticmethod
    def violates_chsh(S: float) -> bool:
        """判断是否违反 CHSH 不等式"""
        return abs(S) > 2.0


# ============================================================================
# Level 6: 路径积分
# ============================================================================

class PathIntegral:
    """
    Feynman 路径积分
    
    ⟨x_f|e^(-iĤt/ℏ)|x_i⟩ = ∫ 𝒟x(t) e^(iS[x]/ℏ)
    
    基于 ±0 的理解:
    - 所有路径的叠加（正负相位干涉）
    - 经典路径：相位极值（作用量极值）
    """
    
    def __init__(self,
                 action_functional: Callable[[np.ndarray], float],
                 hbar: float = QuantumConstants.HBAR_AU):
        """
        Args:
            action_functional: 作用量泛函 S[x(t)]
            hbar: ℏ
        """
        self.S = action_functional
        self.hbar = hbar
    
    def propagator(self,
                  x_initial: float,
                  x_final: float,
                  time: float,
                  num_paths: int = 1000) -> complex:
        """
        传播子（简化 Monte Carlo 计算）
        
        K(x_f, t; x_i, 0) = ⟨x_f|Û(t)|x_i⟩
        """
        # 生成随机路径（简化）
        amplitude = 0.0 + 0.0j
        
        for _ in range(num_paths):
            # 随机路径（线性插值 + 扰动）
            path = self._generate_random_path(x_initial, x_final, time)
            
            # 计算作用量
            action = self.S(path)
            
            # 相位因子
            phase = np.exp(1j * action / self.hbar)
            
            amplitude += phase
        
        return amplitude / num_paths
    
    def _generate_random_path(self,
                             x_i: float,
                             x_f: float,
                             T: float,
                             num_points: int = 10) -> np.ndarray:
        """生成随机路径"""
        t_vals = np.linspace(0, T, num_points)
        
        # 线性插值 + 高斯扰动
        classical_path = x_i + (x_f - x_i) * t_vals / T
        fluctuation = np.random.randn(num_points) * 0.1
        
        return classical_path + fluctuation


# ============================================================================
# 综合演示系统
# ============================================================================

class QuantumDemo:
    """量子力学综合演示"""
    
    @staticmethod
    def demonstrate_wavefunctions():
        """演示波函数"""
        print("\n" + "=" * 70)
        print("波函数与测量")
        print("=" * 70)
        
        # 高斯波包
        domain = (np.array([-10]), np.array([10]))
        psi = WaveFunction.gaussian_wavepacket(x0=0, k0=2, sigma=1, domain=domain)
        psi = psi.normalize()
        
        print("\n高斯波包:")
        print(f"  归一化: {psi.is_normalized()}")
        
        # 期望值
        x_avg = psi.position_expectation()
        p_avg = psi.momentum_expectation()
        
        print(f"  位置期望值 ⟨x⟩ = {x_avg:.4f}")
        print(f"  动量期望值 ⟨p⟩ = {p_avg:.4f}")
    
    @staticmethod
    def demonstrate_operators():
        """演示算符"""
        print("\n" + "=" * 70)
        print("量子算符与不确定性关系")
        print("=" * 70)
        
        dim = 10
        
        # 位置和动量算符
        x_op = PositionOperator(dim, (-5, 5))
        p_op = MomentumOperator(dim, dx=1.0, hbar=1.0)
        
        # 创建态
        state = StateVector.superposition(
            [StateVector.basis_state(0, dim), StateVector.basis_state(1, dim)],
            [1/np.sqrt(2), 1/np.sqrt(2)]
        )
        
        # 不确定性关系
        dx, dp, bound = x_op.uncertainty_relation(p_op, state)
        
        print(f"\nHeisenberg 不确定性关系:")
        print(f"  Δx = {dx:.4f}")
        print(f"  Δp = {dp:.4f}")
        print(f"  Δx·Δp = {dx*dp:.4f}")
        print(f"  下界 ℏ/2 = {bound:.4f}")
        print(f"  满足: {dx*dp >= bound - 1e-10}")
    
    @staticmethod
    def demonstrate_harmonic_oscillator():
        """演示谐振子"""
        print("\n" + "=" * 70)
        print("量子谐振子")
        print("=" * 70)
        
        dim = 5
        omega = 1.0
        H = HamiltonianOperator.harmonic_oscillator(dim, omega)
        
        print(f"\n谐振子 (ω={omega}):")
        print(f"能级:")
        
        for n in range(dim):
            E_n, psi_n = H.energy_eigenstate(n)
            print(f"  n={n}: E={E_n:.4f} ℏω")
        
        # 时间演化
        initial = StateVector.basis_state(0, dim)
        schrodinger = SchrodingerEquation(H)
        
        t_vals, states = schrodinger.time_evolution(initial, (0, 10), num_points=50)
        
        print(f"\n时间演化（基态）:")
        print(f"  初态归一化: {initial.is_normalized()}")
        print(f"  末态归一化: {states[-1].is_normalized()}")
    
    @staticmethod
    def demonstrate_entanglement():
        """演示量子纠缠"""
        print("\n" + "=" * 70)
        print("量子纠缠")
        print("=" * 70)
        
        # Bell 态
        bell = EntangledState.bell_state('phi_plus')
        
        print("\nBell 态 |Φ⁺⟩:")
        print(f"  是否纠缠: {bell.is_entangled()}")
        
        # Schmidt 分解
        schmidt_coeffs, _, _ = bell.schmidt_decomposition()
        print(f"  Schmidt 数: {bell.schmidt_number()}")
        print(f"  Schmidt 系数: {schmidt_coeffs}")
        
        # 纠缠熵
        S_E = bell.entanglement_entropy()
        print(f"  纠缠熵: {S_E:.4f}")
        print(f"  最大纠缠: {abs(S_E - np.log(2)) < 1e-10}")
    
    @staticmethod
    def demonstrate_bell_inequality():
        """演示 Bell 不等式"""
        print("\n" + "=" * 70)
        print("Bell 不等式与非定域性")
        print("=" * 70)
        
        # Bell 态
        bell = EntangledState.bell_state('phi_plus')
        
        # 测量算符（Pauli 矩阵）
        sigma_x = np.array([[0, 1], [1, 0]])
        sigma_z = np.array([[1, 0], [0, -1]])
        
        theta = np.pi / 4
        A1 = sigma_z
        A2 = np.cos(theta) * sigma_z + np.sin(theta) * sigma_x
        B1 = np.cos(theta) * sigma_z - np.sin(theta) * sigma_x
        B2 = sigma_x
        
        # CHSH 参数
        S = BellInequality.chsh_parameter(bell, (A1, A2), (B1, B2))
        
        print(f"\nCHSH 不等式:")
        print(f"  S = {S:.4f}")
        print(f"  经典界限: |S| ≤ 2")
        print(f"  量子界限: |S| ≤ 2√2 ≈ {2*np.sqrt(2):.4f}")
        print(f"  违反经典界限: {BellInequality.violates_chsh(S)}")


def main():
    """主函数：完整演示"""
    
    print("=" * 70)
    print("DaoMath Quantum Mechanics")
    print("从 ±0 构造量子理论")
    print("=" * 70)
    
    # 1. 波函数
    QuantumDemo.demonstrate_wavefunctions()
    
    # 2. 算符与不确定性
    QuantumDemo.demonstrate_operators()
    
    # 3. 谐振子
    QuantumDemo.demonstrate_harmonic_oscillator()
    
    # 4. 量子纠缠
    QuantumDemo.demonstrate_entanglement()
    
    # 5. Bell 不等式
    QuantumDemo.demonstrate_bell_inequality()
    
    print("\n" + "=" * 70)
    print("量子力学构造完成！")
    print("=" * 70)


if __name__ == "__main__":
    main()
